import requests
import numpy as np
from pgvector.django import CosineDistance, L2Distance
from pymupdf.mupdf import vectors
import urllib.parse
from src.modules.knowbase.repository.entity.vectors import TextVector, ImageVector
from sqlalchemy.ext.asyncio import AsyncSession
from src.modules.knowbase.services.Impl.textvector_service import TextVectorService
from src.modules.knowbase.services.Impl.imagevector_service import ImageVectorService
from core.config import settings
from .http_client import HTTPClient
import asyncio
import json
import aiohttp

import logging
logger = logging.getLogger(__name__)

from src.modules.ai_models.model_loaders.model_loader import  get_text_embedding_model


async def retrieve_image_and_text_vectors(know_base_ids, question, session: AsyncSession, top_k=5):
    question_vector = get_embedding(question)
    print(len(question_vector))
    question_vector = question_vector.tolist()
    text_results = await TextVectorService.get_text_vector_list_for_rag(question_vector=question_vector,
                                                                        session=session,
                                                                        knowbase_ids=know_base_ids,
                                                                        top_k=top_k)
    image_results = await ImageVectorService.get_image_vector_list_for_rag(question_vector=question_vector,
                                                                           session=session,
                                                                           knowbase_ids=know_base_ids,
                                                                           top_k=top_k)

    # 将图片路径转换为 URL 可访问的形式
    print(type(text_results))
    print(text_results[0])
    print(text_results)
    print("image_results")
    print(image_results)

    def convert_path_to_url(path):
        normalized_path = path.replace('\\', '/')
        normalized_path= normalized_path.replace('Data','media/uploads')
        normalized_path = "http://localhost:8000/" + normalized_path
        return normalized_path

    for iv, dist in image_results:
        iv.image_path = convert_path_to_url(iv.image_path)
    text_contents = [tv.text_content for tv, dist in text_results]

    # image_paths = [iv.image_path for iv, dist in image_results]
    image_info_list = [(iv.image_desc, iv.image_path) for iv, dist in image_results]
    return text_contents, image_info_list


def retrieve_text_vectors(know_base_ids, qusetion, top_k=5):
    # 取出前 5 个最相关的文本向量
    question_vector = get_embedding(qusetion)
    vectors = TextVector.objects.filter(know_base_id__in=know_base_ids)
    results = vectors.annotate(
        distance=L2Distance('embedding_vector', question_vector)
    ).order_by('distance')[:top_k]
    return [result.embedding_content for result in results]


def retrieve_image_vectors(know_base_ids):
    vectors = ImageVector.objects.filter(know_base_id__in=know_base_ids)
    return [v.image_path for v in vectors[:5]]


async def generate_with_ollama(prompt: str, model: str = None) -> str:
    """使用Ollama API生成文本"""
    model = model or settings.OLLAMA_MODEL
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": True
    }

    session = await HTTPClient.get_session()

    for attempt in range(settings.OLLAMA_MAX_RETRIES):
        try:
            async with session.post(
                    settings.OLLAMA_API_URL,
                    json=payload,
                    raise_for_status=True
            ) as response:
                async for line in response.content:
                    line = line.decode('utf-8').strip()
                    if not line:
                        continue

                    try:
                        data = json.loads(line)
                        if "response" in data:
                            yield data["response"]
                        elif "error" in data:
                            logger.error(f"Ollama API error: {data['error']}")
                            yield f"Error: {data['error']}"
                    except json.JSONDecodeError:
                        logger.warning(f"Failed to parse Ollama response: {line}")
                        continue
                return  # 成功完成，退出函数

        except (aiohttp.ClientError, asyncio.TimeoutError) as e:
            logger.error(f"Attempt {attempt + 1} failed: {str(e)}")
            if attempt < settings.OLLAMA_MAX_RETRIES - 1:
                await asyncio.sleep(2 ** attempt)  # 指数退避
            else:
                yield f"Request to Ollama API failed after {settings.OLLAMA_MAX_RETRIES} attempts: {str(e)}"
                return


def get_embedding(question):
    text_embedding_model = get_text_embedding_model()
    question_embedding = text_embedding_model.get_text_embedding(question)
    return question_embedding



